#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
从Google Drive下载模型权重
"""

import os
import gdown
import argparse

MODEL_URLS = {
    'u2net': 'https://drive.google.com/uc?id=1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ',
    'u2netp': 'https://drive.google.com/uc?id=1IG3HdpcRiDoWNookbncQjeaPN28t90yW',
}

def download_model_weights(model_name):
    """
    下载模型权重文件
    
    Args:
        model_name (str): 模型名称 ('u2net', 'u2netp')
    """
    if model_name not in MODEL_URLS:
        print(f"不支持的模型: {model_name}")
        print(f"支持的模型: {', '.join(MODEL_URLS.keys())}")
        return False
    
    # 确定模型文件路径
    if model_name == 'u2net':
        model_filename = 'u2net_model.pth'
    elif model_name == 'u2netp':
        model_filename = 'u2netp_model.pth'
    else:
        raise ValueError(f"不支持的模型名称: {model_name}")
    
    # 获取模型文件夹路径
    model_dir = os.path.join(os.getcwd(), 'models', 'weights')
    # 确保目录存在
    os.makedirs(model_dir, exist_ok=True)
    
    # 完整的模型文件路径
    model_path = os.path.join(model_dir, model_filename)
    
    # 检查模型是否已经存在
    if os.path.exists(model_path):
        print(f"模型已存在: {model_path}")
        return True
    
    print(f"开始从Google Drive下载模型 {model_name}...")
    url = MODEL_URLS[model_name]
    try:
        gdown.download(url, model_path, quiet=False)
        print(f"模型下载完成: {model_path}")
        return True
    except Exception as e:
        print(f"下载失败: {str(e)}")
        return False

def main():
    parser = argparse.ArgumentParser(description="从Google Drive下载模型权重工具")
    parser.add_argument("model", nargs="?", default="all", 
                        help="要下载的模型名称 (u2net, u2netp, all)")
    
    args = parser.parse_args()
    model_name = args.model.lower()
    
    if model_name == "all":
        print("开始下载所有模型...")
        for model in MODEL_URLS.keys():
            download_model_weights(model)
    elif model_name in MODEL_URLS.keys():
        download_model_weights(model_name)
    else:
        print(f"不支持的模型: {model_name}")
        print(f"支持的模型: {', '.join(MODEL_URLS.keys())}, all")
        return 1
    
    return 0

if __name__ == "__main__":
    main() 